import os
import torch

from typing import Any, Dict, List
from torch.utils.data import Dataset

import random
import matplotlib.pyplot as plt
import numpy as np


class Radar(Dataset):
    def __init__(self, split, cfg: Dict[str, Any]):
        assert split in ['train', 'val'], "split must be 'train' or 'val'"
        self.input_dir = cfg.get('input_path')
        self.label_dir = cfg.get('label_path')
        self.input_files = self._get_pt_files(self.input_dir)
        self.label_files = self._get_pt_files(self.label_dir)
        self.input_size = cfg.get('input_size', [800, 800])
        
        input_names = {os.path.splitext(f)[0]: os.path.join(self.input_dir, f) for f in self.input_files}
        label_names = {os.path.splitext(f)[0]: os.path.join(self.label_dir, f) for f in self.label_files}
        all_names = sorted(list(set(input_names.keys()) & set(label_names.keys())))
        
        num_total = len(all_names)
        num_train = int(num_total * 0.8)
        if split == 'train':
            selected_names = all_names[:num_train]
        else:
            selected_names = all_names[num_train:]
        
        self.pairs = []
        for name in selected_names:
            input_path = input_names[name]
            label_path = label_names[name]
            try:
                input_tensor = torch.load(input_path)
                if isinstance(input_tensor, list) and len(input_tensor) == 8:
                    self.pairs.append((input_path, label_path))
                else:
                    print(f"Warning: {input_path} skipped, input_tensor len={len(input_tensor)}")
            except Exception as e:
                print(f"Warning: {input_path} load error: {e}")

        # self.show_random_label_image()
        # self.plot_transformed_count_distribution()
        # print('s')

    # def show_random_label_image(self):
    #     idx = random.randint(0, len(self)-1)
    #     sample = self[idx]
    #     label = sample['label'].cpu().numpy()  # [H, W]
    #     # 颜色映射：0=蓝色，1=红色，255=灰色
    #     color_map = np.zeros((label.shape[0], label.shape[1], 3), dtype=np.uint8)
    #     color_map[label == 0] = [0, 102, 204]    # blue
    #     color_map[label == 1] = [220, 30, 30]    # red
    #     color_map[label == 255] = [180, 180, 180] # gray

    #     plt.imshow(color_map)
    #     plt.title(f"Label visualization for idx {idx}")
    #     plt.axis("off")
    #     plt.show()
    
    # def plot_transformed_count_distribution(self):
    #     values = []
    #     for i in range(min(10, len(self))):  # 多扫点数据
    #         x = self[i]['x']   # [H, W, 8*16*c]（你合并后是 [H, W, 8*D*c]）
    #         # 按你现在拼接方式，channel 0 就是第一个原始特征的c[0]
    #         feature_map = x[..., 0]  # 所有空间位置的特征0（数量，经过扰动和归一）
    #         valid = feature_map[feature_map != -1]  # 去掉填充
    #         values.append(valid.cpu().numpy())
    #     values = np.concatenate(values)
    #     plt.figure()
    #     plt.hist(values, bins=100)
    #     plt.title(f'After augment/log/normalize: feature 0 ({values.shape[0]} valid values)')
    #     plt.xlabel("Value")
    #     plt.ylabel("Frequency")
    #     plt.show()
    # def plot_input_feature_distribution(self, feature_idx):
    #     """
    #     feature_idx: int, 0~(8*c-1)
    #     """
    #     values = []
    #     for i in range(3):
    #         x = self[i]['x']   # [H, W, 8*c]
    #         feature_map = x[..., feature_idx]
    #         valid = feature_map[feature_map != -1]
    #         values.append(valid.cpu().numpy())
    #     values = np.concatenate(values)
    #     plt.figure()
    #     plt.hist(values, bins=100)
    #     plt.title(f'Feature {feature_idx} distribution ({values.shape[0]} valid values)')
    #     plt.xlabel("Value")
    #     plt.ylabel("Frequency")
    #     plt.show()

    def _get_pt_files(self, folder: str) -> List[str]:
        return sorted([f for f in os.listdir(folder) if f.endswith('.pt')])

    @staticmethod
    def scatter_to_full_tensor(input_tensor, input_size):
        coords, feats = input_tensor
        H, W = input_size
        D = 16
        _, c = feats.shape

        full_tensor = (torch.rand((H, W, D, c), dtype=feats.dtype, device=feats.device) * 2) - 1

        for idx, (x, y, z) in enumerate(coords):
            if 0 <= x < W and 0 <= y < H and 0 <= z < D:
                v = feats[idx].clone()
                v[0] = v[0] + (torch.rand(1, device=v.device) - 0.5) - 0.45
                v[0] = torch.log(torch.clamp(v[0], min=1e-4)) / 3.0
                full_tensor[H-1-y, x, z] = v

        full_tensor = full_tensor.reshape(H, W, D * c)
        return full_tensor

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, index):
        input_path, label_path = self.pairs[index]
        input_tensor = torch.load(input_path)
        label_tensor = torch.load(label_path)
        full_list = []
        for i in range(8):
            coords_and_feats = input_tensor[i]
            full = self.scatter_to_full_tensor(coords_and_feats, self.input_size)
            full_list.append(full)
        input_full = torch.cat(full_list, dim=-1)

        label = label_tensor[0].clone()
        mask = label_tensor[1]

        label[mask == 0] = 255

        return {'x': input_full, 'label': label}
